/**
 * 
 */
package cs534.policy;

import java.util.Arrays;

/**
 * @author jhostetler
 *
 */
public class OnlinePolicyGradient<S, A, BasePolicy extends DifferentiablePolicy<BasePolicy, S, A>> extends MyPolicy<S, A>
{
	private BasePolicy base_;
	private final double[] z_;
	private final double alpha_;
	private final double beta_;
	
	public OnlinePolicyGradient( final BasePolicy base, final double alpha, final double beta )
	{
		base_ = base;
		z_ = new double[base_.parameters().length];
		alpha_ = alpha;
		beta_ = beta;
	}
	
	@Override
	public A getAction()
	{
		System.out.println( s_ );
		final A action = base_.getAction();
		System.out.println( "\n--> Action taken: " + action.toString() );
		return action;
	}
	
	@Override
	public void actionResult( final S s, final A a, final S sprime, final double r )
	{
		// Do a gradient update.
		final double[] grad = base_.gradient( a );
		double grad_norm = 0.0;
		for( int zi = 0; zi < z_.length; ++zi ) {
			z_[zi] = z_[zi]*beta_ + grad[zi];
			grad_norm += (grad[zi]*grad[zi]);
		}
		System.out.println( "Gradient L2 norm = " + Math.sqrt( grad_norm ) );
//		System.out.println( Arrays.toString( z_ ) );
		
		final double[] parameters = Arrays.copyOf( base_.parameters(), base_.parameters().length );
		for( int i = 0; i < parameters.length; ++i ) {
			parameters[i] += alpha_ * r * z_[i];
		}
		
		base_ = base_.update( parameters );
	}
	
	public BasePolicy basePolicy()
	{
		return base_;
	}

	@Override
	protected void enterObservationImpl( final S s )
	{
		base_.enterObservation( s );
	}

}
